from __future__ import annotations
from dataclasses import dataclass, field
from pathlib import Path
from queue import Queue
from typing import Iterable, Optional
from collections import deque
from matplotlib import pyplot as plt
import tqdm
import wave
import numpy as np

_EPS = 1e-8

@dataclass
class TraceGoertzel:
    freq: float  # 目标频率(归一化)
    n: int  # 采样点数(采样率)
    k: Optional[int] = field(default=None)  # int(f * n)
    w: Optional[float] = field(default=None)  # 角频率
    cw: Optional[float] = field(default=None)  # 余弦值
    c: Optional[float] = field(default=None)  # 核心系数C
    sw: Optional[float] = field(default=None)  # 正弦值
    zs: deque = field(
        default_factory=lambda: deque([0.0, 0.0], maxlen=2)
    )  # 存储滤波后的序列 Z[-2, -1]

    def __post_init__(self):
        self.k = int(self.freq * self.n) # 一个周期下目标信号的采样点数
        self.w = 2.0 * np.pi * self.freq # 角频率 w = 2pi * f
        self.cw = np.cos(self.w)
        self.sw = np.sin(self.w)
        self.c = 2 * self.cw
        print(f"Goertzel initialized with freq={self.freq}, n={self.n}, k={self.k}, w={self.w}, c={self.c}, cw={self.cw}, sw={self.sw}")

    def __call__(self, xs: Iterable[float]) -> float:
        """
        Goertzel 核心递归 (二阶数字滤波)
        s[n] = x[n] + 2 * cos(w_k) * s[n-1] - s[n-2]
        """
        self.reset()
        for x in xs:
            z1, z2 = self.zs[-1], self.zs[-2]  # Z[-1], Z[-2]
            z0 = x + self.c * z1 - z2  # S[n] = x[n] + C * S[n-1] - S[n-2]
            self.zs.append(float(z0))  # 更新序列(挤出最左)
        return self.amp  # 返回当前幅度
        
    def reset(self):
        """重置Goertzel状态"""
        self.zs.extend([0.0, 0.0])
    
    @property
    def amp(self) -> float:
        z1, z2 = self.zs[-1], self.zs[-2]
        ip = self.cw * z1 - z2
        qp = self.sw * z1
        return np.sqrt(ip**2 + qp**2) / (self.n / 2.0)

class PairGoertzel:
    def __init__(
        self,
        f_sample: int,
        f_space: int,
        f_mark: int,
        bit_rate: int,  # 比特率, 单位bps, 一位的时长= 1 / bit_rate
        win_size: int,  # 窗口大小, 用于计算goertzel ampitude, 一般 fs/f1 * 10
    ):
        assert f_sample%bit_rate == 0, "采样频率必须是比特率的整数倍"
        self.Fs = f_sample
        self.F0 = f_space
        self.F1 = f_mark
        self.bit_rate = bit_rate
        self.n_per_bit = int(f_sample // bit_rate)  # 每个比特的采样点数
        f1 = f_mark / f_sample  # 归一化频率
        f0 = f_space / f_sample  # 归一化频率
        self.g0 = TraceGoertzel(freq=f0, n=win_size)
        self.g1 = TraceGoertzel(freq=f1, n=win_size)
        self.in_buffer = deque(maxlen=win_size)  # 输入缓冲区, 用于存储采样点
        self.out_buffer = Queue(maxsize=self.n_per_bit) # 信号缓冲区, 用于输出比特
        self.out_count = 0  # 输出计数, 用于控制输出频率
        print(f"Goertzel initialized with f0={f0}, f1={f1}, in_size={win_size}, out_size={self.n_per_bit}")

    def __call__(self, s: float) -> tuple[float, float, float]:
        self.in_buffer.append(s)  # 添加采样点到输入缓冲区
        self.out_count += 1
        amp0, amp1, p1_prob = 0, 0, None
        if self.out_count >= self.n_per_bit:
            amp0 = self.g0(self.in_buffer)  # 对输入缓冲区的采样点进行Goertzel处理
            amp1 = self.g1(self.in_buffer)  # 对输入缓冲区的采样点
            p1_prob = amp1 / (amp0 + amp1 + _EPS)
            self.out_count = 0
        return amp0, amp1, p1_prob

def signal2indicator(
    signal_np: np.ndarray,
    fs: int,
    f0: int,
    f1: int,
    duration: float,
    s_goertzel: int = 5,  # Goertzel窗口大小系数, win = fs/f1 * s_goertzel
):
    goertzel = PairGoertzel(
        f_sample=fs,
        f_space=f0,
        f_mark=f1,
        bit_rate=int(1 / duration),  # 比特率, 单位bps
        win_size=int(fs / f1 * s_goertzel),
    )
    signal_f1, signal_f0 = [], []
    indicators = []
    for s in tqdm.tqdm(signal_np):
        s0, s1, p1_prob = goertzel(s)
        signal_f1.append(float(s1))
        signal_f0.append(float(s0))
        if p1_prob is not None:
            indicators.append(float(p1_prob))

    return signal_f1, signal_f0, indicators

def indicator2bytes(
    indicators: list[float],
    start: bytes = b'\x01\x02',
    end: bytes = b'\x03\x04',
    seperate: bytes = b'\x10',
    threshold: float = 0.5,
):
    signal_bits = ''
    # step1: loop将values转换为bit array
    for v in indicators:
        signal_bits += '1' if v > threshold else '0'

    # step 2: 查找开始符
    start_bits = "".join(format(int(x), '08b') for x in start)
    # 匹配开始符
    start_idx = -1
    for i in range(len(start_bits), len(signal_bits)):
        seq = signal_bits[i-len(start_bits) : i]
        if seq == start_bits:
            start_idx = i
            break
    
    if start_idx == -1:
        print("未找到开始符")
        return []
    
    # step 3: 查找结束符
    end_bits = "".join(format(int(x), '08b') for x in end)
    end_idx = signal_bits.find(end_bits, start_idx)
    
    if end_idx == -1:
        print("未找到结束符")
        return []

    # step 4: 提取有效位
    valid_bits = signal_bits[start_idx:end_idx]
    
    # step 5: 分割有效位为字节
    byte_array = bytearray()
    for i in range(0, len(valid_bits), 8):
        byte_chunk = valid_bits[i:i+8]
        if len(byte_chunk) < 8:
            continue  # 忽略不足8位的部分
        byte_array.append(int(byte_chunk, 2))

    return byte_array.split(seperate)  # 使用分隔符分割字节数组

# %% example
Fs = 6_400  # Sampling frequency (采样频率, 最小也得是载波信号频率的两倍)
F1 = 1800  # mark tone freq
F0 = 1500  # space tone freq
Bitrate = 200
S_GOERTZEL = 9
T = 1 / Bitrate  # Duration per bit (位持续时间)

if __name__ == "__main__":
    # fpath = str((Path(__file__).parent / "temp" /  "test2.wav").resolve())
    # fpath = str((Path(__file__).parent / "temp" / "fsk_signal_1750_1500_25_no_load.wav").resolve())
    # fpath = str((Path(__file__).parent / "temp" / "test4.wav").resolve()) # 18/15/50
    # fpath = str((Path(__file__).parent / "temp" / "test6.wav").resolve()) # 18/15/100
    fpath = str("data/1800_1500_200.wav") # 18/15/200
    with wave.open(fpath, "rb") as wf:
        signal = wf.readframes(wf.getnframes())  # 读取所有帧
        signal = np.frombuffer(signal, dtype=np.float16)
        fs_actual = wf.getframerate()
        print(f"实际采样率: {fs_actual} Hz")
    print("load signal from", fpath)
    # 如果和Fs不一致, 需要重新采样
    if fs_actual != Fs:
        print(f"实际采样率 {fs_actual} 与预期采样率 {Fs} 不一致, 需要重新采样")
        # 重新采样到Fs
        from scipy.signal import resample
        signal = resample(signal, int(len(signal) * Fs / fs_actual))
        fs_actual = Fs

    # 解调信号(Goertzel算法)
    signal_f1, signal_f0, indicators = signal2indicator(
        signal,
        fs=fs_actual,
        f0=F0,
        f1=F1,
        duration=T,
        s_goertzel=S_GOERTZEL
    )

    # 将指示器转换为字节
    results = indicator2bytes(
        indicators,
        start=b'\x01\x02',
        end=b'\x03\x04',
        seperate=b'\x10',
        threshold=0.5
    )

    # 打印解调结果
    print("\n解调结果:")
    for i, r in enumerate(results):
        try:
            decoded_text = r.decode('utf-8')
            print(f"结果 {i}: '{decoded_text}' (原始字节: {r.hex()})")
        except UnicodeDecodeError:
            print(f"结果 {i}: 无法解码为UTF-8 (原始字节: {r.hex()})")

    # 生成时间轴
    duration_total = len(signal) / fs_actual
    time_signal = np.linspace(0, duration_total, len(signal))
    time_goertzel = np.linspace(0, duration_total, len(signal_f1))
    time_indicators = np.linspace(0, duration_total, len(indicators))

    # %% 创建综合图形
    plt.figure(figsize=(12, 6))
    
    # 子图1: 原始信号
    plt.subplot(3, 1, 1)
    plt.plot(time_signal, signal, 'b-', linewidth=0.5, alpha=0.7)
    plt.ylabel('幅度')
    plt.title('原始信号')
    plt.grid(True, alpha=0.3)
    plt.xlim(0, duration_total)
    
    # 子图2: Goertzel幅度检测
    plt.subplot(3, 1, 2)
    plt.plot(time_goertzel, signal_f0, 'g-', linewidth=1, label=f'F0({F0}Hz) 幅度')
    plt.plot(time_goertzel, signal_f1, 'r-', linewidth=1, label=f'F1({F1}Hz) 幅度')
    plt.ylabel('Goertzel幅度')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.xlim(0, duration_total)

    # 子图3: 频谱分析
    plt.subplot(3, 1, 3)
    frequencies = np.fft.fftfreq(len(signal), 1/fs_actual)
    fft_signal = np.abs(np.fft.fft(signal))
    positive_freq_idx = frequencies >= 0
    plt.plot(frequencies[positive_freq_idx], fft_signal[positive_freq_idx], 'orange', linewidth=2)
    plt.axvline(x=F0, color='g', linestyle='--', linewidth=1, label=f'Space频率: {F0}Hz')
    plt.axvline(x=F1, color='r', linestyle='--', linewidth=1, label=f'Mark频率: {F1}Hz')
    plt.ylabel('频谱幅度')
    plt.xlabel('频率 (Hz)')
    plt.title('频谱分析')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.xlim(0, 3000)
    
    plt.tight_layout()
    plt.show(block=False)
    
    # 如果有解调结果，创建比特序列图
    if indicators:
        plt.figure(figsize=(12, 6))
        
        # 转换为比特序列
        bit_sequence = ''.join(['1' if p > 0.5 else '0' for p in indicators])
        
        plt.subplot(2, 1, 1)
        plt.plot(time_indicators, indicators, 'purple', linewidth=2)
        plt.axhline(y=0.5, color='red', linestyle='--', alpha=0.7)
        
        # 背景着色显示比特值
        for i, (t, p) in enumerate(zip(time_indicators, indicators)):
            bit = 1 if p > 0.5 else 0
            color = 'lightblue' if bit == 1 else 'lightgray'
            if i < len(time_indicators) - 1:
                plt.axvspan(t, time_indicators[i+1], alpha=0.3, color=color)
        
        plt.ylabel('P1概率')
        plt.title('比特判决过程')
        plt.grid(True, alpha=0.3)
        plt.ylim(0, 1)
        
        plt.subplot(2, 1, 2)
        # 显示比特序列（前N位，避免过长）
        display_n = 200
        display_bits = bit_sequence[:display_n] if len(bit_sequence) > display_n else bit_sequence
        bit_values = [int(b) for b in display_bits]
        bit_times = np.linspace(0, len(display_bits)/len(bit_sequence) * duration_total, len(display_bits))
        
        plt.step(bit_times, bit_values, 'b-', linewidth=2, where='post')
        plt.ylabel('比特值')
        plt.xlabel('时间 (秒)')
        plt.title(f'解调比特序列 (前{len(display_bits)}位)')
        plt.grid(True, alpha=0.3)
        plt.ylim(-0.5, 1.5)
        
        plt.tight_layout()
        plt.show(block=False)
    
    # 打印统计信息
    print(f"信号长度: {len(signal)} 采样点")
    print(f"信号持续时间: {duration_total:.3f} 秒")
    print(f"解调出的比特数: {len(indicators)}")
    print(f"平均比特率: {len(indicators)/duration_total:.1f} bps")